#pragma once
#include "net_common.hpp"
#include "net_message.hpp"

namespace dw
{
    namespace net
    {
        // type "T" is user-defined enum class,
        // indicating the message type
        template <typename T>
        struct message_header
        {
            union {
                T type;
                uint32_t id;
            };
            uint32_t size;  // body size in bytes

        public:
            message_header()
            {
                id = 0;
                size = 0;
            }
        
        public:
            message_header& NetOrder(){
                this->id = htonl(this->id);
                this->size = htonl(this->size);
                
                return *this;
            }

            message_header& HostOrder()
            {
                this->type = T(ntohl(this->id));
                this->size = ntohl(this->size);

                return *this;
            }

        public:
            friend std::ostream& operator << (std::ostream& os, const message_header& mh)
            {
                os << "Type ID: " << mh.id << " Size: " << mh.size << " Bytes";
                return os;
            }
        };

        template <typename T>
        struct message
        {
            message_header<T> header;
            std::vector<uint8_t> body;
        
        public:
            message()
            {
                // Empty Constructor
                // done
            }
            message(T _type)
            {
                header.type = _type;
                header.size = 0;
            }

            ~message()
            {
                // done
            }

        public:
            // Message size: include header size
            size_t size() const
            {
                return sizeof(message_header<T>) + body.size();
            }

            T type() const
            {
                return header.type;
            }

            bool checkType(T _type)
            {
                return _type == header.type;
            }

            void setType(T _type)
            {
                header.type = _type;
            }

        public:
            message& append(const std::string &str)
            {
                uint32_t i = this->body.size();
                uint32_t len = uint32_t(str.size());

                this->body.resize(this->body.size() + len);
                std::memcpy((void *)(this->body.data() + i), (const void *)str.data(), len);

                this->header.size = this->body.size();

                return *this;
            }

            std::string data()
            {
                std::string str;

                uint32_t len = this->body.size();
                // pop out raw string data
                str.resize(len);
                std::memcpy((void *)str.data(), (const void *)(this->body.data()), len);

                return str;
            }

            message& clear()
            {
                this->body.clear();
                return *this;
            }

        public:
            /* Specialize string input and output*/

            friend message<T> &operator << (message<T> &msg, const std::string &str)
            {
                uint32_t i = msg.body.size();
                uint32_t len = uint32_t(str.size());

                msg.body.resize(msg.body.size() + len + sizeof(uint32_t));
				// push in raw string data
				std::memcpy((void*)(msg.body.data() + i), (const void*)str.data(), len);
				// push in string data length in Network Endian
				len = htonl(len);
				std::memcpy((void*)(msg.body.data() + i + str.size()), (const void*)&len, sizeof(uint32_t));

				msg.header.size = msg.body.size();

				return msg;
            }

            friend message<T>& operator >> (message<T>& msg, std::string &str)
            {
                uint32_t i = msg.body.size() - sizeof(uint32_t);

                // pop out string data length in Network Endian
                uint32_t len;
                std::memcpy((void *)&len, (const void *)(msg.body.data() + i), sizeof(uint32_t));
                len = ntohl(len);

                // pop out raw string data
                str.resize(len);
                std::memcpy((void *)str.data(), (const void *)(msg.body.data() + i - len), len);

                msg.body.resize(i - len);
                msg.header.size = msg.body.size();

                return msg;
            }

            /* template for trivial data type input and output*/

            template <typename DataType>
            friend message<T> &operator<<(message<T> &msg, const DataType &data)
            {
                // Check that the type of the data being pushed is trivially copyable
                static_assert(std::is_standard_layout<DataType>::value, "Data is too complex to be pushed into vector!");

                uint32_t i = msg.body.size();
                msg.body.resize(msg.body.size() + sizeof(DataType));
                std::memcpy((void *)(msg.body.data() + i), (const void *)&data, sizeof(DataType));

                msg.header.size = msg.body.size();

                return msg;
            }

            template <typename DataType>
            friend message<T> &operator>>(message<T> &msg, const DataType &data)
            {
                // Check that the type of the data being pushed is trivially copyable
                static_assert(std::is_standard_layout<DataType>::value, "Data is too complex to be extracted from vector!");

                uint32_t i = msg.body.size() - sizeof(DataType);
                std::memcpy((void *)&data, (const void *)(msg.body.data() + i), sizeof(DataType));
                msg.body.resize(i);

                msg.header.size = msg.body.size();

                return msg;
            }

        public:
            // Output Message Type ID and body content in hex
            friend std::ostream &operator<<(std::ostream &os, const message<T> &msg)
            {
                os << msg.header << std::endl;

                uint32_t cnt = 0;
                os << std::hex;
                for (auto &byte : msg.body)
                {
                    switch (cnt % 20)
                    {
                        case 0:
                            os << std::endl;
                            break;
                        default:
                            os << " ";
                            if (0 == cnt % 4) {
                                os << " ";
                            }
                            break;
                    }
                    cnt++;
                    os << std::setw(2) << std::setfill('0') << (uint32_t)byte;
                }
                os << std::dec;
                return os;
            }

        };

        template <typename T>
        class connection;

        template <typename T>
        struct owned_message
        {   
            // Used in server_interface and its connections, 
            // to determine which client sent this messsage

            std::shared_ptr<connection<T>> remote = nullptr;
            message<T> msg;

            friend std::ostream &operator<<(std::ostream &os, const owned_message<T> &msg)
            {
                if (msg.remote) {
                    os << "Connection ID: " << msg.remote->GetID() << std::endl;
                }
                os << msg.msg;
                return os;
            }
        };
    }
}